# sequence_data_collection.py
import cv2
import mediapipe as mp
import numpy as np
import os
# import time
from collections import deque

# 初始化MediaPipe
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(static_image_mode=False, max_num_hands=1, min_detection_confidence=0.7)
mp_drawing = mp.solutions.drawing_utils

# 定义手势序列类别
SEQUENCES = ['swipe_down', 'swipe_up', 'grab']
current_sequence = 0
is_recording = False
frame_buffer = deque(maxlen=10)  # 存储连续帧的关键点
recorded_sequences = 0

# 创建数据保存目录
if not os.path.exists('sequence_data'):
    os.makedirs('sequence_data')
    for seq in SEQUENCES:
        os.makedirs(f'sequence_data/{seq}')


def extract_landmarks(results):
    """从MediaPipe结果中提取关键点坐标"""
    if results.multi_hand_landmarks:
        landmarks = []
        for landmark in results.multi_hand_landmarks[0].landmark:
            landmarks.extend([landmark.x, landmark.y, landmark.z])
        return np.array(landmarks)
    return None


cap = cv2.VideoCapture(0)

print("手势序列数据采集")
print("按数字键选择手势序列：1-向下滑动, 2-向上滑动, 3-抓取")
print("按空格键开始/结束录制序列，按q退出")
print("提示：每个序列持续约1-2秒")

while True:
    ret, frame = cap.read()
    if not ret:
        break

    frame = cv2.flip(frame, 1)
    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 手势检测
    results = hands.process(rgb_frame)

    # 显示当前状态
    status_text = f"Sequence: {SEQUENCES[current_sequence]} | Recorded: {recorded_sequences}"
    if is_recording:
        status_text += " | RECORDING"

    cv2.putText(frame, status_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

    if results.multi_hand_landmarks:
        for hand_landmarks in results.multi_hand_landmarks:
            mp_drawing.draw_landmarks(frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)

            # 如果正在录制，保存关键点序列
            if is_recording:
                landmarks = extract_landmarks(results)
                if landmarks is not None:
                    frame_buffer.append(landmarks)

    cv2.imshow('Sequence Data Collection', frame)

    key = cv2.waitKey(1) & 0xFF

    if key == ord('q'):
        break
    elif key in [ord('1'), ord('2'), ord('3')]:
        current_sequence = key - ord('1')
        recorded_sequences = 0
        print(f"切换到手势序列: {SEQUENCES[current_sequence]}")
    elif key == ord(' '):  # 空格键开始/结束录制
        if not is_recording:
            # 开始录制
            is_recording = True
            frame_buffer.clear()
            print(f"开始录制 {SEQUENCES[current_sequence]}...")
        else:
            # 结束录制并保存
            is_recording = False
            if len(frame_buffer) >= 5:  # 至少5帧才保存
                # 将序列保存为单个npy文件
                sequence_data = np.array(frame_buffer)
                save_path = f'sequence_data/{SEQUENCES[current_sequence]}/seq_{recorded_sequences:04d}.npy'
                np.save(save_path, sequence_data)
                recorded_sequences += 1
                print(f"已保存 {SEQUENCES[current_sequence]} 序列 #{recorded_sequences}, 帧数: {len(frame_buffer)}")
            else:
                print("序列太短，未保存")

cap.release()
cv2.destroyAllWindows()